import asyncio
import json
from openai import OpenAI
from fastmcp import FastMCP, Client
from typing import List, Dict, Any

# --- 0. 配置 ---
# 文件名
TRAINING_DATA_FILE = "batch_dataset.jsonl" 

# 配置 OpenAI 客户端 (请替换您的 API 密钥和 Base URL)
client = OpenAI(
    base_url="https://ai.gitee.com/v1",
    api_key="XXXXXXXXX", 
    default_headers={"X-Failover-Enabled":"true"},
)

# --- 1. FastMCP 多个工具定义和客户端 ---

mcp = FastMCP('demo.mcp')

@mcp.tool()
def get_current_weather(city: str) -> str:
    """
    查询指定城市当前的实时天气。
    输入参数：
    city: 要查询的城市名称，例如“北京”或“伦敦”。这个参数是必需的，并且应该是字符串类型。
    输出参数：
    result: 格式化的天气描述。
    """
    # 按照您的要求，始终返回相同的结果
    return f"{city} 的天气是多云，温度 35 度。"

@mcp.tool()
def calculate_sum(a: float, b: float) -> float:
    """
    计算两个数字的和。
    输入参数：
    a: 第一个数字，必需。
    b: 第二个数字，必需。
    """
    return a + b

mcp_client = Client(mcp)

async def execute_tool(tool_name: str, tool_args: dict) -> str:
    """实际异步调用 FastMCP 工具，提取结果数据，并返回 JSON 字符串。"""
    print(f"\n[执行工具: {tool_name}] 参数: {tool_args}")
    async with mcp_client:
        try:
            result_obj = await mcp_client.call_tool(tool_name, tool_args)
            
            actual_data = result_obj.data 
            
            print(f"✅ 工具执行完成，原始数据: {actual_data}")
            
            return json.dumps({"result": actual_data}) 
        
        except Exception as e:
            print(f"❌ 工具执行错误: {e}")
            return json.dumps({"error": str(e)})


# --- 2. 多个工具的 JSON Schema 定义 ---

WEATHER_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_current_weather",
        "description": "查询指定城市当前的实时天气。",
        "parameters": {
            "type": "object",
            "properties": {"city": {"type": "string", "description": "要查询的城市名称，例如“北京”或“伦敦”。"} },
            "required": ["city"]
        }
    }
}

CALCULATE_SCHEMA = {
    "type": "function",
    "function": {
        "name": "calculate_sum",
        "description": "计算两个数字的和。",
        "parameters": {
            "type": "object",
            "properties": {
                "a": {"type": "number", "description": "第一个数字。"},
                "b": {"type": "number", "description": "第二个数字。"}
            },
            "required": ["a", "b"]
        }
    }
}

# 组合所有可用的工具 Schema
AVAILABLE_TOOLS = [WEATHER_SCHEMA, CALCULATE_SCHEMA]


# --- 3. 数据保存函数 (保持不变) ---

def save_training_data(messages: List[Dict[str, Any]], filename: str):
    """将一次完整的对话历史保存为 JSONL 格式的训练数据。"""
    data = {"messages": messages}
    try:
        with open(filename, 'a', encoding='utf-8') as f:
            json_line = json.dumps(data, ensure_ascii=False)
            f.write(json_line + '\n')
        print(f"\n✅ 对话数据已成功保存到 {filename}\n{'='*50}")
    except Exception as e:
        print(f"❌ 保存训练数据失败: {e}")


# --- 4. 核心异步 Function Calling 逻辑 ---

async def handle_single_conversation(user_query: str):
    """处理单个用户查询的完整 Agent 流程。"""
    print(f"\n--- 🚀 开始处理新的查询：'{user_query[:30]}...' ---")

    # 初始对话消息 (每次循环都从头开始)
    messages = [
        {
            "role": "system",
            "content": f"你是一个善于使用工具的助手。如果用户的请求可以用提供的工具解决，请调用工具。可用的工具列表：\n{json.dumps(AVAILABLE_TOOLS, ensure_ascii=False)}"
        },
        {
            "role": "user",
            "content": user_query
        }
    ]

    print("--- 1. 第一次 LLM 调用：判断意图并选择工具 ---")
    
    response = client.chat.completions.create(
        model="DeepSeek-R1",
        messages=messages,
        tools=AVAILABLE_TOOLS, 
        tool_choice="auto",
        max_tokens=8192,
    )

    response_message = response.choices[0].message
    
    # **修复点：在保存模型消息前，使用 .model_dump() 转换为字典**
    messages.append(response_message.model_dump()) 
    print(f"模型响应: {'（Tool Call）' if response_message.tool_calls else response_message.content}")

    # --- Tool Execution Logic ---
    if response_message.tool_calls:
        tool_calls = response_message.tool_calls
        tool_call = tool_calls[0]
        tool_name = tool_call.function.name
        tool_args = json.loads(tool_call.function.arguments)
        
        print(f"🤖 模型选择调用工具：**{tool_name}**")

        tool_result_content = await execute_tool(tool_name, tool_args)
        
        # 将工具调用的结果添加到对话历史中
        messages.append(
            {
                "tool_call_id": tool_call.id,
                "role": "tool",
                "name": tool_name,
                "content": tool_result_content,
            }
        )

        # --- 2. 第二次 LLM 调用：生成最终回复 ---
        print("\n--- 2. 第二次 LLM 调用：生成最终回复 ---")
        
        final_response = client.chat.completions.create(
            model="DeepSeek-R1",
            messages=messages,
            max_tokens=8192,
        )

        final_response_message = final_response.choices[0].message
        final_content = final_response_message.content
        
        # **修复点：在保存最终回复前，使用 .model_dump() 转换为字典**
        messages.append(final_response_message.model_dump())
        
        print(f"🎉 最终回复:\n{final_content}")
        
    else:
        # 模型未调用工具，直接回复的消息已在前面添加
        print(f"模型未调用工具，直接回复:\n{response_message.content}")

    # --- 5. 保存数据 ---
    save_training_data(messages, TRAINING_DATA_FILE)


async def run_batch_conversations():
    # ----------------------------------------------------
    # 您的 3 个不同用户查询
    # ----------------------------------------------------
    CONVERSATIONS = [
        # 样本 1: 触发 get_current_weather 工具
        "请问**上海**今天的天气情况怎么样？",
        
        # 样本 2: 触发 calculate_sum 工具
        "我需要知道 42 和 100 的和是多少，请用你的计算工具。",
        
        # 样本 3: 不触发任何工具
        "中国的首都在哪里",
    ]
    
    # 确保文件是空的，以便从头开始记录
    open(TRAINING_DATA_FILE, 'w').close()
    
    for query in CONVERSATIONS:
        await handle_single_conversation(query)
        # 延迟一下，避免 API 速率限制
        await asyncio.sleep(1) 


# 这是推荐的、唯一的程序入口点
if __name__ == '__main__':
    # 运行批处理主函数
    asyncio.run(run_batch_conversations())